package com.github.common.extend.demo.util.excel;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.support.ExcelTypeEnum;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.alibaba.excel.write.style.column.LongestMatchColumnWidthStyleStrategy;
import com.github.common.extend.demo.util.XRunnable;
import com.github.common.extend.demo.util.excel.dynamic.DynamicTitleDefine;
import com.github.common.extend.demo.util.excel.dynamic.DynamicTitleReadListener;
import com.github.common.extend.demo.util.excel.dynamic.DynamicTitleRow;
import com.github.common.extend.demo.util.excel.dynamic.DynamicTitleWriteHandler;
import com.github.common.extend.demo.util.iterator.XIterableUtils;
import com.github.common.extend.demo.util.iterator.XProducerIterator;
import com.github.common.extend.demo.util.iterator.page.XPageIterable;
import com.github.common.extend.demo.util.iterator.page.XPageParam;
import com.github.common.extend.demo.util.iterator.partition.XPartitionFunction;
import com.github.common.extend.demo.util.iterator.partition.XPartitionIterator;
import com.github.common.extend.demo.util.iterator.partition.fixed.LocalDateTimePartitionFunction;
import com.github.common.extend.demo.util.iterator.partition.fixed.LocalDateTimePartitionParam;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import java.io.InputStream;
import java.io.OutputStream;
import java.time.LocalDateTime;
import java.time.temporal.TemporalUnit;
import java.util.*;
import java.util.concurrent.*;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * excel写工具
 * @author wangjj7
 * @date 2023/7/27
 * @description
 */
@Slf4j
public class XExcelUtils {

    /**
     * 单个sheet保存的数据默认上限
     */
    public static final int DEFAULT_SHEET_MAX_COUNT = 1000000;

    /**
     * 每次从iterator中取默认数量的元素写入
     */
    public static final int DEFAULT_SINGLE_WRITE_ROW_COUNT = 10000;

    /**
     * 生成excel文件并写入流
     * <p>
     * 优化点
     * 1.单个sheet限制最大数据写入数量
     * 超过限制数量自动写入下一个sheet
     * 2.支持分批加载（需要以懒加载的方式实现迭代器）
     * 避免一次读取大量数据导致内存溢出
     * 3.支持读写并行（需要实现分片加载函数且配置异步分片读取执行器，当然写入还是串行的）
     * 方法调用线程写入当前分片数据时，执行器线程提前并行读取下一分片数据
     *
     * 假设分片数为N，每个分片读取耗时为R，写入耗时为W，方法调用线程为M ，异步执行器工作线程为T，同步执行总耗时【（R + W） * N】，
     * 则每个分区平均节省的耗时为【R + W - MAX(R ,W)】；总节约的耗时为【MIN(R ,W) + MAX(R ,W) * N】
     *
     * 场景分析：
     * 1.当 R ≈ W，线程执行总耗时【(N + 1) * (R or W)】，线程阻塞总耗时【0】，接近同步写入的一半（推荐场景）；
     * |---R(M)---|---W(M)---|
     *            |---R(T)---|---W(M)---|
     *                       |---R(T)---|---W(M)---|
     *                                  |---R(T)---|---W(M)---|
     *                                             |---R(T)---|---W(M)---|
     * M end
     *
     * 2.当 R << W，线程执行总耗时【R + N * W】，线程执行阻塞总耗时【0】；
     * |-R(M)-|---W(M)---|
     *        |-R(T)-|   |---W(M)---|
     *                   |-R(T)-|   |---W(M)---|
     *                              |-R(T)-|   |---W(M)---|
     *                                         |-R(T)-|   |---W(M)---|
     * M end
     *
     * 3.当 W << R，线程执行总耗时【N * R + W】，线程（M）执行阻塞总耗时【(R - W) * (N - 1)】；调用线程M频繁阻塞
     * |---R(M)---|-W(M)-|
     *            |---R(T)---|-W(M)-|
     *                       |---R(T)---|-W(M)-|
     *                                  |---R(T)---|-W(M)-|
     *                                             |---R(T)---|-W(M)-|
     * M end
     * <p>
     *
     * @param sheetMaxCount 每个sheet最大写入数量
     * @param excelWriter
     * @param writeSheetFunction 根据sheetNo生成WriteSheet
     * @param iterator  元素迭代器
     * @param singleWriteRowCount 每次从iterator中取指定数量的元素写入
     *                  配合懒加载方法实现的iterator（例如数据库分页查询），如果是Collection的话使用默认值即可
     * @param readExecutor 异步分片读取执行器；如果为null或无可执行线程，则同步读取
     *                  配合懒加载方法实现的iterator（例如数据库分页查询），如果是Collection的话则没有必要设置
     *
     * @return excel已经写入完成
     */
    public static <T> void write(int sheetMaxCount, ExcelWriter excelWriter,
                                 Function<Integer ,WriteSheet> writeSheetFunction,
                                 Iterator<T> iterator, int singleWriteRowCount, Executor readExecutor)
    {
        if (sheetMaxCount <= 0 || iterator == null || singleWriteRowCount < 1) {
            throw new IllegalArgumentException();
        }

        int sheetNo = 0;
        int sheetRowCount = 0;

        WriteSheet writeSheet = writeSheetFunction.apply(sheetNo);
        //第一个分片由当前线程直接加载
        List<T> partitionList = XIterableUtils.next(iterator, singleWriteRowCount);

        while (true) {
            if (CollectionUtils.isEmpty(partitionList)){
                //结束任务
                excelWriter.write(Collections.EMPTY_LIST, writeSheet);
                excelWriter.finish();
                return;
            }

            int needWriteCount = partitionList.size();
            int canWriteCount = sheetMaxCount - sheetRowCount;

            if (canWriteCount >= needWriteCount) {
                //当前sheet可以全部写入
                if (!iterator.hasNext()) {
                    //写入后结束任务
                    excelWriter.write(partitionList, writeSheet);
                    excelWriter.finish();
                    return;
                }
                XRunnable<List<T>> xRunnable = new XRunnable<>(
                        () -> XIterableUtils.next(iterator ,singleWriteRowCount)
                );
                if (readExecutor != null) {
                    //当前线程写入excel的时候下一个数据区间可以通过异步线程并行读取
                    readExecutor.execute(xRunnable);
                }
                //写入
                excelWriter.write(partitionList, writeSheet);
                sheetRowCount += partitionList.size();

                //如果异步线程还没有开始读取，由当前线程执行读取任务,否则等待读取结果
                xRunnable.run();
                try {
                    partitionList = xRunnable.getFuture().get();
                } catch (ExecutionException e) {
                    throw new IllegalStateException(e);
                } catch (InterruptedException e) {
                    throw new IllegalStateException(e);
                }
            } else {
                if (canWriteCount > 0) {
                    excelWriter.write(partitionList.subList(0, canWriteCount), writeSheet);
                }
                //将溢出的数据写到下一个sheet里
                partitionList = partitionList.subList(canWriteCount, needWriteCount);
                writeSheet = EasyExcel.writerSheet(++sheetNo).build();
                sheetRowCount = 0;
                continue;
            }
        }
    }

    /**
     * 重载
     * @see #write(int, ExcelWriter, Function, Iterator, int, Executor)
     *
     * @param sheetMaxCount     每个sheet最大写入数量
     * @param headType          表头类型
     * @param outputStream      excel生成后写入的流
     * @param iterator 元素迭代器
     * @param singleWriteRowCount 每次从iterator中取指定数量的元素写入
     *                            配合懒加载方法实现的iterator（例如数据库分页查询），如果是Collection的话使用默认值即可
     * @param readExecutor      异步分片读取执行器；如果为null或无可执行线程，则同步读取
     *                          配合懒加载方法实现的iterator（例如数据库分页查询），如果是Collection的话则没有必要设置
     * @return excel已经写入完成
     */
    public static <T> void write(int sheetMaxCount, Class<T> headType, OutputStream outputStream,
                                 Iterator<T> iterator, int singleWriteRowCount ,Executor readExecutor)
    {
        if (sheetMaxCount <= 0 || headType == null || outputStream == null
                || iterator == null || singleWriteRowCount < 1) {
            throw new IllegalArgumentException();
        }

        try (ExcelWriter excelWriter = EasyExcel.write(outputStream, headType).excelType(ExcelTypeEnum.XLSX).build()) {
            write(sheetMaxCount ,excelWriter ,sheetNo -> EasyExcel.writerSheet(sheetNo).build() ,
                    iterator ,singleWriteRowCount ,readExecutor);
        }
    }

    /**
     * 重载
     * @see #write(int, Class, OutputStream, Iterator, int, Executor)
     *
     * @param headType          表头类型
     * @param outputStream      excel生成后写入的流
     * @param iterator 元素迭代器
     * @param readExecutor      异步分片读取执行器；如果为null或无可执行线程，则同步读取
     *                          配合懒加载方法实现的iterator（例如数据库分页查询），如果是Collection的话则没有必要设置
     * @return excel已经写入完成
     */
    public static <T> void write(Class<T> headType, OutputStream outputStream, Iterator<T> iterator, Executor readExecutor)
    {
        write(DEFAULT_SHEET_MAX_COUNT ,headType ,outputStream ,iterator ,DEFAULT_SINGLE_WRITE_ROW_COUNT ,readExecutor);
    }

    /**
     * 根据分片函数的实现批次写入
     *
     * 重载
     * @see #write(Class, OutputStream, Iterator, Executor)
     *
     * @param headType
     * @param outputStream
     * @param partitionFunction
     * @param readExecutor
     * @param <T>
     * @param <P>
     */
    public static <T, P> void writeByPartition(Class<T> headType, OutputStream outputStream,
                                               XPartitionFunction<T, P> partitionFunction, Executor readExecutor)
    {
        write(headType ,outputStream ,new XPartitionIterator<>(partitionFunction) ,readExecutor);
    }

    /**
     * 生成excel文件并写入流（异步）
     *
     * @param headType
     * @param outputStream
     * @param iterator
     * @param writeExecutor 异步写入执行器
     * @param readExecutor  异步分片读取执行器；如果为null或无可执行线程，则同步读取
     *                      配合懒加载方法实现的iterator（例如数据库分页查询），如果是Collection的话则没有必要设置
     * @param <T>
     */
    public static <T> Future<Void> asyncWrite(Class<T> headType, OutputStream outputStream,
                                              Iterator<T> iterator, Executor writeExecutor,
                                              Executor readExecutor)
    {
        if (headType == null || outputStream == null || iterator == null
                || writeExecutor == null || readExecutor == null)
        {
            throw new IllegalArgumentException();
        }

        CompletableFuture<Void> future = new CompletableFuture<>();
        writeExecutor.execute(() -> {
            write(headType ,outputStream ,iterator ,readExecutor);
            future.complete(null);
        });
        return future;
    }

    /**
     * 生成excel文件并写入流（异步）
     *
     * 重载
     * @see #asyncWrite(Class, OutputStream, Iterator, Executor, Executor)
     *
     * @param headType
     * @param outputStream
     * @param iterator
     * @param executor 如果需要支持并行读写，可执行线程数至少有两个（一个写，一个读），如果只有一个可执行线程则读写串行串行
     * @param <T>
     * @param <P>
     */
    public static <T ,P> Future<Void> asyncWrite(Class<T> headType, OutputStream outputStream,
                                                 Iterator<T> iterator, Executor executor)
    {
        return asyncWrite(headType, outputStream, iterator, executor ,executor);
    }

    /**
     * 根据时间分片的规则分批将数据写入excel
     * 如果结果集很大且查询参数带时间范围，可以将时间拆分成一个个小时间分片后通过此方法分批写入
     *
     * @param outputStream 生成的excel写入的流
     * @param query 查询条件
     * @param headType excel头类型
     * @param partitionSize 时间分片大小
     * @param partitionInterval 时间分片间隔
     * @param intervalUnit 分片单位
     * @param asc true:升序，false:降序
     * @param queryFunction 分片查询函数
     * @param readExecutor  异步分片读取执行器；如果为null或无可执行线程，则同步读取
     *                      配合懒加载方法实现的iterator（例如数据库分页查询），如果是Collection的话则没有必要设置
     *
     * @param <T>
     * @param <P>
     */
    public static <T ,P extends LocalDateTimePartitionParam<P>> void writeByTimePartition(
            OutputStream outputStream, P query, Class<T> headType,
            int partitionSize, int partitionInterval, TemporalUnit intervalUnit, boolean asc,
            Function<P, List<T>> queryFunction , ExecutorService readExecutor)
    {
        LocalDateTime startLocalTime = query.$readStartTime();
        LocalDateTime endLocalTime = query.$readEndTime();

        if (startLocalTime == null || endLocalTime == null){
            throw new IllegalArgumentException("时间范围必须指定");
        }
        if (partitionSize < 1){
            throw new IllegalArgumentException("分区大小必须大于0");
        }

        XPartitionFunction<T, P> partitionFunction = null;
        if (startLocalTime.isAfter(endLocalTime)){
            partitionFunction = input -> new XPartitionFunction.Output<T ,P>().setHasNext(false);
        }else {
            partitionFunction = new LocalDateTimePartitionFunction<>(query ,
                    partitionSize ,intervalUnit ,p -> queryFunction.apply(p))
                    .setAsc(asc)
                    .setPartitionInterval(partitionInterval);
        }

        XExcelUtils.writeByPartition(headType, outputStream, partitionFunction, readExecutor);
    }

    /**
     * 同步阻塞写入excel
     *
     * 重载
     * @see #write(Class, OutputStream, Iterator, Executor)
     *
     * @param headType
     * @param outputStream
     * @param iterator
     * @param <T>
     */
    public static <T> void write(Class<T> headType, OutputStream outputStream, Iterator<T> iterator)
    {
        write(headType, outputStream, iterator, null);
    }

    /**
     * 同步阻塞写入excel
     *
     * 重载
     * @see #write(Class, OutputStream, Iterator)
     *
     * @param headType
     * @param outputStream
     * @param iterable
     * @param <T>
     */
    public static <T> void write(Class<T> headType, OutputStream outputStream, Iterable<T> iterable)
    {
        write(headType, outputStream, iterable.iterator());
    }

    /**
     * 带自定义标题的excel写入
     * @param dynamicTitleDefines
     * @param headType
     * @param rowIterable
     * @param outputStream
     * @param <T>
     */
    public static <K ,T extends DynamicTitleRow<K>> void writeForDynamicTitle(
            Class<T> headType, List<DynamicTitleDefine<K>> dynamicTitleDefines,
            Iterable<T> rowIterable, OutputStream outputStream)
    {
        if (headType == null || outputStream == null || rowIterable == null) {
            throw new IllegalArgumentException();
        }

        //如果迭代过程中涉及查询，直接生成两个iterator会导致重复查询。既然DynamicTitleWriteHandler中迭代的元素都是每次写入的元素，保存每次迭代的元素即可
        XProducerIterator producerIterator = XIterableUtils.subscribe(rowIterable.iterator());

        try (ExcelWriter excelWriter = EasyExcel.write(outputStream, headType).excelType(ExcelTypeEnum.XLSX).build()) {
            write(DEFAULT_SHEET_MAX_COUNT ,excelWriter ,sheetNo -> EasyExcel.writerSheet(sheetNo)
                            .registerWriteHandler(new DynamicTitleWriteHandler(dynamicTitleDefines, producerIterator.consumer()))
                            .registerWriteHandler(new LongestMatchColumnWidthStyleStrategy())
                            .build() ,
                    producerIterator ,DEFAULT_SINGLE_WRITE_ROW_COUNT ,null);
        }
    }

    /**
     * 分页查询分批写入带自定义标题的excel
     *
     * 重载
     * @see #writeForDynamicTitle(Class, List, Iterable, OutputStream)
     *
     * @param headType
     * @param dynamicTitleDefines
     * @param outputStream
     * @param pageQueryFunction
     * @param pageSize
     * @param <T>
     */
    public static <K ,T extends DynamicTitleRow<K>> void writePageForDynamicTitle(
            Class<T> headType, List<DynamicTitleDefine<K>> dynamicTitleDefines ,OutputStream outputStream,
            Function<XPageParam, List<T>> pageQueryFunction, int pageSize)
    {
        XPageIterable xPageIterable = new XPageIterable(param -> pageQueryFunction.apply(param) ,pageSize);
        writeForDynamicTitle(headType ,dynamicTitleDefines ,xPageIterable ,outputStream);
    }

    /**
     * 分批加载excel
     * @param inputStream
     * @param partitionCount 每批加载的行数
     * @param headType
     * @param resultCallback 每加载一批数据后执行的操作
     * @param <T>
     */
    public static <T> void read(InputStream inputStream ,int partitionCount, Class<T> headType,
                                   Consumer<List<T>> resultCallback)
    {
        EasyExcel.read(inputStream)
                .head(headType)
                //分批加载写入
                .registerReadListener(new PartitionFlushReadListener(partitionCount, resultCallback))
                .sheet()
                .doRead();
    }

    /**
     * 分批加载带自定义标题的excel
     * @param inputStream
     * @param dynamicTitleDefines 自定义标题定义
     * @param headType
     * @param resultCallback 每加载一批数据后执行的操作
     * @param <K>
     * @param <T>
     */
    public static <K ,T extends DynamicTitleRow<K>> void readForDynamicTitle(
            InputStream inputStream, int partitionCount, Class<T> headType,
            List<DynamicTitleDefine<K>> dynamicTitleDefines, Consumer<List<T>> resultCallback)
    {
        EasyExcel.read(inputStream)
                .head(headType)
                //自定义标签处理
                .registerReadListener(new DynamicTitleReadListener<>(dynamicTitleDefines))
                //分批加载写入
                .registerReadListener(new PartitionFlushReadListener(partitionCount, resultCallback))
                .sheet()
                .doRead();
    }

}
